from sklearn.decomposition import PCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
import numpy as np

def pca(traces,point_num,*args):
    pca = PCA(n_components=point_num)
    return pca.fit_transform(traces)

def lda(traces,point_num,midvalue,train_range,*args):
    if len(midvalue.shape) == 2:
        midvalue = midvalue @ [2 ** (7 - i) for i in range(8)]
    lda = LinearDiscriminantAnalysis(n_components=point_num)
    traces_train=lda.fit_transform(traces[train_range[0]:train_range[1]], midvalue)
    traces_test=lda.transform(np.vstack((traces[:train_range[0]],traces[train_range[1]:])))
    return np.vstack((traces_test[:train_range[0]],traces_train,traces_test[train_range[0]:]))

def poi_corrcoef(traces,point_num,midvalue,*args):
    if len(midvalue.shape)==2:
        midvalue = midvalue @ [2 ** (7 - i) for i in range(8)]
    corrcoef_point=np.abs(np.corrcoef(midvalue,traces.transpose())[0,1:])
    corrcoef_sort = np.argsort(-corrcoef_point)
    return corrcoef_sort[:point_num]

def poi_snr(traces, point_num,*args):
    traces_mean = np.mean(traces, axis=0)
    traces_std = np.std(traces, axis=0, ddof=0)
    SNR = np.abs(np.where(traces_std == 0, 0, traces_mean / traces_std))
    SNR_sort = np.argsort(SNR)[::-1]
    return SNR_sort[:point_num]